import requests
from typing import Union


class MCPClient:
    def __init__(self, base_url: str = "http://localhost:8000"):
        self.base_url = base_url

    def calculate_squares(self, a: Union[int, float], b: Union[int, float]) -> dict:
        """
        调用平方和计算接口
        :param a: 第一个数字
        :param b: 第二个数字
        :return: 包含计算结果的字典
        """
        response = requests.post(
            f"{self.base_url}/square_sum",
            json={"a": a, "b": b}  # 使用JSON body传参
        )
        response.raise_for_status()  # 如果请求失败抛出异常
        return response.json()

    def get_greeting(self, name: str) -> dict:
        """
        调用问候接口
        :param name: 要问候的名字
        :return: 包含问候语的字典
        """
        response = requests.get(
            f"{self.base_url}/resources/greeting/{name}"  # 路径参数传参
        )
        response.raise_for_status()
        return response.json()


if __name__ == "__main__":
    # 创建客户端实例
    client = MCPClient()

    try:
        # 测试平方和计算接口
        print("测试平方和计算:")
        result = client.calculate_squares(3.5, 2.5)
        print(f"3.5² + 2.5² = {result['sum']} (详情: {result})")

        # 测试问候接口
        print("\n测试问候接口:")
        greeting = client.get_greeting("Alice")
        print(greeting["result"])

    except requests.exceptions.RequestException as e:
        print(f"请求失败: {e}")
    except KeyError as e:
        print(f"响应数据格式错误，缺少字段: {e}")